#include <stdio.h>          // 标准输入输出头文件，用于printf等函数
#include <cuda_runtime.h>   // CUDA运行时库头文件，提供GPU相关函数
#include <chrono>           // C++时间库，用于高精度计时
#include <cstdio>           // C标准输入输出，用于sprintf等函数
#include <ratio>            // 用于表示时间单位比例，如秒、毫秒
#include <string>           // 字符串库，用于处理字符串
#include <iostream>         // 输入输出流，用于cout等

// 初始化矩阵函数
// 参数：data-矩阵数据指针；size-矩阵大小；min-最小值；max-最大值；seed-随机种子
void initMatrix(float* data, int size, int min, int max, int seed) {
    srand(seed);  // 设置随机种子
    for (int i = 0; i < size; i ++) {
        // 生成[min, max)范围内的随机浮点数（注意：当前实现未正确包含min，应为min + ...）
        data[i] = float(rand()) * float(max - min) / RAND_MAX; 
    }
}

// 打印矩阵数据
// 参数：data-矩阵数据指针；size-矩阵大小
void printMat(float* data, int size) {
    for (int i = 0; i < size; i ++) {
        printf("%.8lf", data[i]); // 打印每个元素，保留8位小数
        if (i != size - 1) {
            printf(", ");
        } else {
            printf("\n");
        }
    }
}

// 比较两个矩阵是否相同（允许1e-4的误差）
// 参数：h_data-CPU结果指针；d_data-GPU结果指针；size-矩阵大小
void compareMat(float* h_data, float* d_data, int size) {
    double precision = 1.0E-4; // 允许的误差精度
    bool error = false;
    for (int i = 0; i < size; i ++) {
        // 检查每个元素的绝对误差
        if (abs(h_data[i] - d_data[i]) > precision) {
            error = true;
            printf("结果在%d处不同, CPU: %.8lf, GPU: %.8lf\n",i, h_data[i], d_data[i]);
            break; // 发现错误即停止检查
        }
    }
    if (error) 
      printf("矩阵乘法结果不同\n");
    else
      printf("矩阵乘法结果相同，误差精度1.0E-4\n");
}

// 计时器类，用于测量代码执行时间
class Timer {
public:
    using s  = std::ratio<1, 1>;    // 秒
    using ms = std::ratio<1, 1000>;  // 毫秒
    using us = std::ratio<1, 1000000>; // 微秒
    using ns = std::ratio<1, 1000000000>; // 纳秒

public:
    Timer(){}; // 构造函数

    void start() { mStart = std::chrono::high_resolution_clock::now(); } // 开始计时
    void stop()  { mStop  = std::chrono::high_resolution_clock::now(); } // 结束计时

    // 计算持续时间并输出结果，span为时间单位，msg为输出信息
    template <typename span>
    void duration(std::string msg);

private:
    std::chrono::time_point<std::chrono::high_resolution_clock> mStart; // 开始时间点
    std::chrono::time_point<std::chrono::high_resolution_clock> mStop;  // 结束时间点
};

// 实现duration方法
template <typename span>
void Timer::duration(std::string msg){
    std::string str; // 时间单位字符串
    char fMsg[100];  // 格式化消息缓冲区
    std::sprintf(fMsg, "%-30s", msg.c_str()); // 格式化消息左对齐

    // 根据时间单位类型设置字符串
    if(std::is_same<span, s>::value) { str = " s"; }
    else if(std::is_same<span, ms>::value) { str = " ms"; }
    else if(std::is_same<span, us>::value) { str = " us"; }
    else if(std::is_same<span, ns>::value) { str = " ns"; }

    // 计算时间差并输出
    std::chrono::duration<double, span> time = mStop - mStart;
    std::cout << fMsg << " 耗时 " << time.count() << str << std::endl;
}

// CUDA核函数：矩阵乘法
// 每个线程计算P的一个元素
__global__ void MatmulKernel(float *M_device, float *N_device, float *P_device, int width){
    // 计算当前线程处理的全局坐标（x,y）
    int x = blockIdx.x * blockDim.x + threadIdx.x; // 列索引
    int y = blockIdx.y * blockDim.y + threadIdx.y; // 行索引

    float P_element = 0; // 存储累加结果

    // 遍历矩阵的行和列进行乘加运算
    for (int k = 0; k < width; k ++){
        float M_element = M_device[y * width + k]; // M的第y行第k列元素
        float N_element = N_device[k * width + x]; // N的第k行第x列元素
        P_element += M_element * N_element;       // 累加到结果
    }

    P_device[y * width + x] = P_element; // 将结果写入P的对应位置
}

// GPU矩阵乘法函数
// 参数：M_host,N_host-输入矩阵；P_host-结果矩阵；width-矩阵宽度；blockSize-线程块大小
void MatmulOnDevice(float *M_host, float *N_host, float* P_host, int width, int blockSize){
    int size = width * width * sizeof(float); // 矩阵总字节数

    // 分配设备内存
    float *M_device, *N_device;
    cudaMalloc(&M_device, size); // 分配设备端M矩阵内存
    cudaMalloc(&N_device, size); // 分配设备端N矩阵内存

    // 将数据从主机拷贝到设备
    cudaMemcpy(M_device, M_host, size, cudaMemcpyHostToDevice);
    cudaMemcpy(N_device, N_host, size, cudaMemcpyHostToDevice);

    // 分配设备端结果矩阵内存
    float *P_device;
    cudaMalloc(&P_device, size);

    // 设置核函数执行参数
    dim3 dimBlock(blockSize, blockSize); // 线程块维度（blockSize x blockSize）
    dim3 dimGrid(width / blockSize, width / blockSize); // 网格维度
    
    // 调用核函数，注意：实际应检查cudaPeekAtLastError()和cudaDeviceSynchronize()
    MatmulKernel <<<dimGrid, dimBlock>>> (M_device, N_device, P_device, width);

    // 将结果从设备拷贝回主机
    cudaMemcpy(P_host, P_device, size, cudaMemcpyDeviceToHost);
    cudaDeviceSynchronize(); // 同步设备，确保所有操作完成

    // 释放设备内存
    cudaFree(P_device);
    cudaFree(N_device);
    cudaFree(M_device);
}

// CPU矩阵乘法函数（三重循环实现）
void MatmulOnHost(float *M, float *N, float *P, int width){
    for (int i = 0; i < width; i ++)        // 遍历P的行
        for (int j = 0; j < width; j ++){   // 遍历P的列
            float sum = 0;
            for (int k = 0; k < width; k++) // 计算每个元素的内积
                sum += M[i * width + k] * N[k * width + j];
            P[i * width + j] = sum;
        }
}



int seed;
int main(){
    Timer timer; // 创建计时器
    int width     = 1<<10; // 矩阵宽度1024（2^10）
    int min       = 0;     // 随机数最小值
    int max       = 1;     // 随机数最大值
    int size      = width * width; // 矩阵元素总数
    int blockSize = 1;     // 初始块大小（后续测试不同值）

    // 分配主机内存
    //3. 为什么需要 malloc？
    //动态内存分配：程序运行时才决定需要多少内存（比如根据用户输入或文件大小）。
    //避免栈溢出：大内存（如大型数组）不宜在栈上分配（可能导致崩溃），堆空间更大。
    //灵活控制生命周期：手动分配的内存会一直存在，直到调用 free() 释放。
    //malloc 是 C 语言中动态管理内存的核心工具，它让你能灵活控制程序的内存使用。记住：
    //分配内存 → 检查是否成功 → 使用内存 → 释放内存。
    //动态内存的生命周期完全由程序员控制，务必谨慎操作！



    float* h_matM = (float*)malloc(size * sizeof(float));
    float* h_matN = (float*)malloc(size * sizeof(float));
    float* h_matP = (float*)malloc(size * sizeof(float)); // CPU结果
    float* d_matP = (float*)malloc(size * sizeof(float)); // GPU结果

    // 初始化输入矩阵（注意：当前实现中min未被正确应用）
    seed = 1;
    initMatrix(h_matM, size, min, max, seed);
    seed += 1;
    initMatrix(h_matN, size, min, max, seed);

    // CPU端计算并计时
    timer.start();
    MatmulOnHost(h_matM, h_matN, h_matP, width);
    timer.stop();
    timer.duration<Timer::ms>("CPU矩阵乘法");

    // GPU预热（首次运行可能有额外开销）
    timer.start();
    MatmulOnDevice(h_matM, h_matN, d_matP, width, blockSize);
    timer.stop();
    timer.duration<Timer::ms>("GPU矩阵乘法（预热）");

    // 测试不同块大小对性能的影响
    blockSize = 16; // 块大小16x16
    timer.start();
    MatmulOnDevice(h_matM, h_matN, d_matP, width, blockSize);
    timer.stop();
    timer.duration<Timer::ms>("GPU矩阵乘法（块大小16）");
    compareMat(h_matP, d_matP, size); // 验证结果正确性

    blockSize = 1; // 块大小1x1
    timer.start();
    MatmulOnDevice(h_matM, h_matN, d_matP, width, blockSize);
    timer.stop();
    timer.duration<Timer::ms>("GPU矩阵乘法（块大小1）");
    compareMat(h_matP, d_matP, size);

    blockSize = 32; // 块大小32x32
    timer.start();
    MatmulOnDevice(h_matM, h_matN, d_matP, width, blockSize);
    timer.stop();
    timer.duration<Timer::ms>("GPU矩阵乘法（块大小32）");
    compareMat(h_matP, d_matP, size);

    // 释放主机内存
    free(h_matM);
    free(h_matN);
    free(h_matP);
    free(d_matP);

    return 0;
}